/* uc65 - Micro Code for 6502/NES
 * Copyright (C) 2013  Norman B. Lancaster
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.gmail.qbradq.uc65;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Calendar;
import java.util.EnumSet;
import java.util.Set;

/**
 * The class responsible for generating output assembly code.
 */
public class CodeGenerator {
	private static final Set<Operation> readOps = EnumSet.of(
		Operation.ADC,
		Operation.AND,
		Operation.ASL,
		Operation.EOR,
		Operation.CMP,
		Operation.LDA,
		Operation.LDX,
		Operation.LDY,
		Operation.LSR,
		Operation.ORA,
		Operation.ROL,
		Operation.ROR,
		Operation.SBC,
		Operation.INC,
		Operation.DEC
	);
	private static final Set<Operation> writeOps = EnumSet.of(
		Operation.STA,
		Operation.INC,
		Operation.DEC
	);
	
	/**
	 * @param op The opcode to check.
	 * @return True if the opcode is a read operation, false otherwise.
	 */
	public static boolean isReadOp(Operation op) {
		return readOps.contains(op);
	}
	
	/**
	 * @param op The opcode to check.
	 * @return True if the opcode is a write operation, false otherwise.
	 */
	public static boolean isWriteOp(Operation op) {
		return writeOps.contains(op);
	}
	
	private Uc65 compiler;
	private File outFile;
	private StringBuilder header = new StringBuilder();
	private StringBuilder body = new StringBuilder();
	private StringBuilder exports = new StringBuilder();
	private StringBuilder imports = new StringBuilder();
	private StringBuilder out;
	private boolean haveReportedWriteError = false;
	private boolean inRam = false;
	private boolean inRom = false;
	private int ramBank = -1;
	private int romBank = -1;
	private int nextLabel = 0;
	private Optimizer optimizer;
	
	/**
	 * Constructs a new CodeGenerator for a given source file.
	 * @param path The path of the source file.
	 * @param compiler The compiler to report errors to.
	 */
	public CodeGenerator(File source, Uc65 compiler) {
		optimizer = new Optimizer();
		this.compiler = compiler;
		imports.append(String.format("; Assembly generated by uc65\n; " +
				"File: %s\n", source.getPath()));
		imports.append(String.format("; Built Time: %s\n",
			Calendar.getInstance().getTime()));
		if(Uc65.emitDebugInfo) {
			imports.append(String.format(".DBG FILE, \"%s\", %d, %d\n",
				source.getPath(), source.length(),
				(int)(source.lastModified() / 1000L)));
		}
		imports.append("\n\n");
		header.append("\n\n");
		body.append("\n\n");
		exports.append("\n\n");
		out = body;
	}
	
	/** Selects which portion of the file to write to.
	 * 
	 * @param header If true, selects the header section. Otherwise selects the
	 *        body section.
	 */
	public void selectSourceSection(boolean header) {
		out = header ? this.header : this.body;
	}
	
	/** Emits an import statement for a label.
	 * 
	 * @param label The label name
	 * @param zeroPage If true the label will be forced to zero-page.
	 */
	public void emitImport(String label, boolean zeroPage) {
		if(zeroPage) {
			imports.append(String.format(".IMPORTZP\t%s\n", label));
		} else {
			imports.append(String.format(".IMPORT\t\t%s\n", label));
		}
	}
	
	/**
	 * @param name The name to transform.
	 * @return A label-safe string for the given name.
	 */
	public String labelNameFrom(String name) {
		return name.replaceAll("[^a-zA-Z0-9_]", "_");
	}
	
	/**
	 * Closes the CodeGenerator and flushes output to disk.
	 */
	public void close() {
		try {
			outFile = new File(Uc65.outPath);
			BufferedWriter outWriter;
			try {
				outWriter = new BufferedWriter(new FileWriter(outFile));
			} catch(IOException e) {
				reportWriteError();
				return;
			}
			outWriter.write(imports.toString());
			outWriter.write(header.toString());
			outWriter.write(body.toString());
			outWriter.write(exports.toString());
			outWriter.flush();
			outWriter.close();
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * @return An assembler label unique to this translation unit.
	 */
	public String label() {
		return String.format("_L%06X", nextLabel++);
	}
	
	private void reportWriteError() {
		if(!haveReportedWriteError) {
			compiler.error(String.format("Unable to write to output file %s",
				outFile.getPath()), null);
			haveReportedWriteError = true;
		}
	}
	
	private void switchRomBank(int bank) {
		if(!inRom || bank != romBank) {
			inRom = true;
			inRam = false;
			romBank = bank;
			out.append(String.format(".SEGMENT \"ROM%d\"\n", bank));
		}
	}
	
	private void switchRamBank(int bank) {
		if(!inRam || bank != ramBank) {
			inRom = false;
			inRam = true;
			ramBank = bank;
			if(bank == 0) {
				out.append(String.format(".SEGMENT \"BSS%d\": zeropage\n",
					bank));
			} else {
				out.append(String.format(".SEGMENT \"BSS%d\"\n", bank));				
			}
		}
	}
	
	/**
	 * Emits a label to the assembly output.
	 * @param name The label identifier.
	 */
	public void emitLabel(int romBank, String name) {
		switchRomBank(romBank);
		out.append(String.format("\t%s:\n", name));
	}
	
	/**
	 * Emits a variable label.
	 * @param line The source line that defined this variable.
	 * @param name The name of the variable.
	 * @param length The length of the variable, may be zero. Only valid if
	 *        readOnly is false.
	 * @param extern Flag to say if the label should be exported.
	 * @param bank The memory bank to place this variable in.
	 * @param readOnly If true place the variable in the program store,
	 *        otherwise place it in the BSS.
	 * @param data Initial data to be placed in the program store, only valid
	 *        if readOnly is true.
	 */
	public void emitVariable(SourceLine line, String name, int length,
		boolean extern, int bank, boolean readOnly, byte[] data) {
		if(readOnly) {
			switchRomBank(bank);
			out.append(String.format("\t%s:", name));
			for(int i = 0; i < data.length; ++i) {
				String hex = Integer.toHexString(0x100 | data[i]);
				hex = hex.substring(hex.length() - 2);
				if((i % 12) == 0) {
					out.append("\n\t\t.byte $" + hex);
				} else {
					out.append(", $" + hex);						
				}
			}
			out.append("\n");
		} else {
			switchRamBank(bank);
			out.append(String.format("\t%s: .RES %d\n", name, length));
		}
		if(extern) {
			if(bank == 0 && !readOnly) {
				exports.append(String.format(".EXPORTZP\t%s\n", name));				
			} else {			
				exports.append(String.format(".EXPORT\t\t%s\n", name));
			}
		}
		if(Uc65.emitDebugInfo) {
			out.append(String.format(
				".DBG SYM, \"%s\", \"00\", %s, \"%s\"\n",
				name, extern ? "EXTERN" : "STATIC", name));
		}
	}
	
	/**
	 * Enters an assembly procedure block.
	 * @param line The source line the procedure was defined on.
	 * @param name The name of the procedure.
	 * @param extern Flag to say if the label should be exported.
	 * @param romBank The rom bank to place this procedure.
	 */
	public void enterProcedure(SourceLine line, String name, boolean extern,
		int romBank) {
		switchRomBank(romBank);
		emitSourceLine(line);
		out.append(String.format(".PROC %s\n", name));
		if(extern) {
			exports.append(String.format(".EXPORT\t\t%s\n", name));
		}
		if(Uc65.emitDebugInfo) {
			out.append(String.format(
				".DBG FUNC, \"%s\", \"00\", %s, \"%s\"\n",
				name, extern ? "EXTERN" : "STATIC", name));
		}
	}
	
	/**
	 * Exits an assembly procedure block.
	 */
	public void exitProcedure() {
		out.append(".ENDPROC\n");
	}
	
	/**
	 * Emits a source code line as a comment.
	 * @param line The source code line to emit.
	 */
	public void emitSourceLine(SourceLine line) {
		if(Uc65.emitSourceLines) {
			out.append(String.format(";%% %s\n", line.getText()));
		}
		if(Uc65.emitDebugInfo) {
			out.append(String.format(".DBG LINE, \"%s\", %d\n",
				line.getPath(), line.getLine()));
		}
	}
	
	/**
	 * Emits a source line to the output file verbatim.
	 * @param line The line to emit.
	 */
	public void emitVerbatimLine(SourceLine line) {
		emitSourceLine(line);
		out.append(String.format("\t%s\n", line.getText()));
	}
	
	/**
	 * Emits an "assign-style" label to a given address.
	 * @param name The name of the label.
	 * @param value The address to assign the label to.
	 */
	public void emitAssign(String name, int value) {
		out.append(String.format("%s := $%04X\n", name, value));
	}
	
	/**
	 * Called when a new, non-deterministic execution block begins.
	 */
	public void newExecutionBlock() {
		optimizer.reset();
	}
	
	/**
	 * Emits a single operation.
	 * @param op The opcode to emit.
	 * @param mode The addressing mode to use.
	 * @param romBank The ROM bank in which to place code.
	 * @param name An assembly label to use for addressed modes.
	 * @param value A byte value for immediate addressing.
	 * @param vol A flag to indicate if the instruction is volatile.
	 */
	public void emitOpCode(Operation op, AddressMode mode, int romBank,
		String name, byte value, boolean vol) {
		Optimizer.Instruction isr = Optimizer.Instruction.create(
			op, mode, name, value);
		Optimizer.Instruction optimal = optimizer.instruction(isr);
		
		// If the instruction was not volatile, replace it with what the
		// optimizer suggests. If the optimizer says we can omit the
		// the instruction, just return.
		if(!vol) {
			isr = optimal;
		}
		if(isr == null) {
			return;
		}
		op = isr.op;
		mode = isr.mode;
		name = isr.ref;
		value = isr.value;
		
		String code = op.toString().toLowerCase();
		switchRomBank(romBank);
		switch(mode) {
			case IMPLICIT:
				out.append(String.format("\t%s\n", code));
				break;
			case IMMEDIATE:
				out.append(String.format("\t%s\t#$%02X\n", code, value));
				break;
			case ADDRESS:
				out.append(String.format("\t%s\t%s\n", code, name));
				break;
			case ADDRESSX:
				out.append(String.format("\t%s\t%s,x\n", code, name));
				break;
			case ADDRESSY:
				out.append(String.format("\t%s\t%s,y\n", code, name));
				break;
			case INDIRECT:
				out.append(String.format("\t%s\t(%s)\n", code, name));
				break;
			case INDIRECTX:
				out.append(String.format("\t%s\t(%s,x)\n", code, name));
				break;
			case INDIRECTY:
				out.append(String.format("\t%s\t(%s),y\n", code, name));
				break;
			case LABELLO:
				out.append(String.format("\t%s\t#<%s\n", code, name));
				break;
			case LABELHI:
				out.append(String.format("\t%s\t#>%s\n", code, name));
				break;
		}
	}
}
